# !/usr/bin/python
# _*_ coding: utf-8 _*_
# @Time : 2021/5/25 3:21 下午
# @Author : liujun
# @File : base.py
# @desc : 公用CRUD方法

from typing import TypeVar, Generic, Type, Optional, Any, Dict, List
from pydantic import BaseModel
from sqlalchemy import desc, asc
from sqlalchemy.orm import Session

from fastapi.encoders import jsonable_encoder

from db.base import Base

ModelType = TypeVar('ModelType', bound=Base)
AddSchemaType = TypeVar('AddSchemaType', bound=BaseModel)
UpdSchemaType = TypeVar('UpdSchemaType', bound=BaseModel)


class Base(Generic[ModelType, AddSchemaType, UpdSchemaType]):

    def __init__(self, model: Type[ModelType]):
        self.model = model

    def get_filters(self, dic_args: dict):
        filters = []
        for key, value in dic_args.items():
            if key.endswith('_gt'):
                key = key[:-3]
                filters.append(getattr(self.model, key) > value)
            elif key.endswith('_lt'):
                key = key[:-3]
                filters.append(getattr(self.model, key) < value)
            elif key.endswith('_ge'):
                key = key[:-3]
                filters.append(getattr(self.model, key) >= value)
            elif key.endswith('_le'):
                key = key[:-3]
                filters.append(getattr(self.model, key) <= value)
            elif key.endswith('_like'):
                key = key[:-5]
                filters.append(getattr(self.model, key).like('%'+value+'%'))
            elif key.endswith('_in'):
                key = key[:-3]
                filters.append(getattr(self.model, key).in_(value))
            elif key.endswith('_nin'):
                key = key[:-4]
                filters.append(getattr(self.model, key).notin_(value))
            elif key.endswith('_ne'):
                key = key[:-3]
                filters.append(getattr(self.model, key) != value)
            else:
                filters.append(getattr(self.model, key) == value)
        return filters

    def get(self, db: Session, id: int) -> Optional[ModelType]:
        db_obj = db.query(self.model).get(id)
        return jsonable_encoder(db_obj)

    def retrieve(
            self,
            db: Session,
            filters: Dict = dict(),
            sort_name: str = None,
            sort_order: str = None,
            page_number: int = 1,
            page_size: int = 10
    ) -> Optional[List[ModelType]]:
        offset = page_size * (page_number - 1)
        query = db.query(self.model)
        if filters:
            query = query.filter(*self.get_filters(filters))
        if sort_name is not None and sort_order is not None:
            if sort_order == 'desc':
                query = query.order_by(desc(getattr(self.model, sort_name)))
            else:
                query = query.order_by(asc(getattr(self.model, sort_name)))
        query = query.offset(offset).limit(page_size)
        return query.all()

    def count(self, db: Session, filters: Dict = dict()) -> Any:
        query = db.query(self.model)
        if filters:
            query = query.filter(*self.get_filters(filters))
        return query.count()

    def create(self, db: Session, obj_in: AddSchemaType) -> Any:
        obj_in_data = jsonable_encoder(obj_in)
        db_obj = self.model(**obj_in_data)
        db.add(db_obj)
        db.commit()
        db.refresh(db_obj)
        return jsonable_encoder(db_obj)

    def update(self, db: Session, id: int, obj_in: [UpdSchemaType, Dict[str, Any]]) -> Any:
        db_obj = db.query(self.model).get(id)
        obj_data = db_obj.as_dict()
        if isinstance(obj_in, dict):
            update_data = obj_in
        else:
            update_data = obj_in.dict(exclude_unset=True)
        for field in obj_data:
            if field in update_data:
                setattr(db_obj, field, update_data[field])
        db.commit()
        db.refresh(db_obj)
        return jsonable_encoder(db_obj)

    def delete(self, db: Session, ids) -> Any:
        if isinstance(ids, int):
            obj = db.query(self.model).get(ids)
            db.delete(obj)
            db.commit()
        elif isinstance(ids, list):
            for i in ids:
                obj = db.query(self.model).get(i)
                db.delete(obj)
            db.commit()
        else:
            raise ValueError
        return

